package world.snowcrystal.nova.ai;

import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import com.zhipu.oapi.service.v4.model.ModelData;
import io.reactivex.Flowable;
import org.springframework.stereotype.Component;
import world.snowcrystal.nova.ai.Entity.ChatCompletionRequestFactory;

@Component
public class AiTemplateImpl implements AITemplate {

    private final ClientV4 clientV4;

    private final ChatCompletionRequestFactory chatCompletionRequestFactory;


    public AiTemplateImpl(ClientV4 clientV4, ChatCompletionRequestFactory chatCompletionRequestFactory) {
        this.clientV4 = clientV4;
        this.chatCompletionRequestFactory = chatCompletionRequestFactory;
    }


    @Override
    public String chat(String message, boolean stream, float temperature) {
        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(chatCompletionRequestFactory.create(message, stream, temperature));
        return modelApiResponse.getData().getChoices().get(0).toString();
    }

    @Override
    public String chat(String systemMessage, String userMessage, boolean stream, float temperature) {
        ModelApiResponse modelApiResponse = clientV4
                .invokeModelApi(chatCompletionRequestFactory.create(systemMessage, userMessage, stream, temperature));
        return modelApiResponse
                .getData()
                .getChoices()
                .get(0)
                .getMessage()
                .getContent()
                .toString();
    }

    @Override
    public Flowable<ModelData> chatStream(String message, float temperature) {

        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(chatCompletionRequestFactory.createSSE(message, temperature));

        return modelApiResponse.getFlowable();
    }

    @Override
    public Flowable<ModelData> chatStream(String systemMessage, String message, float temperature) {
        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(chatCompletionRequestFactory.createSSE(systemMessage, message, temperature));
        return modelApiResponse.getFlowable();
    }
}
